//
// Created by yunnan xu on 2021/6/6.
//

#include "SkinnedMeshModel.h"
#include "MeshModel.h"

#include <utility>
#include <sstream>
#include "assimp/scene.h"
#include "assimp/anim.h"
#include <gtc/type_ptr.hpp>
#include <gtx/matrix_decompose.hpp>
#include <gtx/quaternion.hpp>

// # define SKINMESHED_DEBUG

bool BoneVertexData_sorter(BoneVertexData const& lhs, BoneVertexData const& rhs) {
    return lhs.bone_weight > rhs.bone_weight;
}

SkinnedMeshModel::SkinnedMeshModel(const char *modelPath, std::unordered_set<VertexAttribType> attribList):
    MeshModel(modelPath, std::move(attribList))
{
    aiMesh* first_mesh = meshes[0];
    if (!first_mesh->HasBones())
    {
        LOGI("SkinnedMesh Has No bones, %s", modelPath);
        return;
    }

    LOGI("SkinnedMesh anim number = %d", ai_scene->mNumAnimations);

    std::vector<std::vector<BoneVertexData>> vertex_weight;
    vertex_weight.resize(vertexNumber);

    for (int i = 0; i < ai_scene->mNumAnimations; i++)
    {
        auto anim = ai_scene->mAnimations[i];
#ifdef SKINMESHED_DEBUG
        LOGI("SkinnedMesh anim name = %s, mDuration = %f, mTicksPerSecond = %f, mNumChannels = %d, mNumMeshChannels = %d",
                anim->mName.C_Str(),
                anim->mDuration,
                anim->mTicksPerSecond,
                anim->mNumChannels,
                anim->mNumMeshChannels);
#endif

        for (int j = 0; j < anim->mNumChannels; j++)
        {
            auto node_anim = anim->mChannels[j];
// #ifdef SKINMESHED_DEBUG

            LOGI("SkinnedMesh node_anim %s, mNumPositionKeys = %d, mNumRotationKeys = %d, mNumScalingKeys = %d",
                    node_anim->mNodeName.C_Str(),
                    node_anim->mNumPositionKeys,
                    node_anim->mNumRotationKeys,
                    node_anim->mNumScalingKeys);

#ifdef SKINMESHED_DEBUG
            std::stringstream ss;
            for (int k = 0; k < node_anim->mNumPositionKeys; k ++)
            {
                ss << node_anim->mPositionKeys[k].mTime << ", ";
            }

            LOGI("SkinnedMesh mPositionKeys = %s", ss.str().c_str());
#endif

// #endif
        }

        for (int j = 0; j < anim->mNumMeshChannels; j ++)
        {
            auto vertex_anim = anim->mMeshChannels[j];
// #ifdef SKINMESHED_DEBUG
            LOGI("SkinnedMesh vertex_anim %s, mNumKeys = %d", vertex_anim->mName.C_Str(), vertex_anim->mNumKeys);
// #endif
        }
    }

    LOGI("SkinnedMesh bones number = %d", first_mesh->mNumBones);

    InitVertexAttrib(VertexAttribType::BONE1_I4);
    InitVertexAttrib(VertexAttribType::WEIGHT1_F4);

    int * bone_index_data = (int *) GetVertexData(VertexAttribType::BONE1_I4);
    float * bone_weight_data = (float *) GetVertexData(VertexAttribType::WEIGHT1_F4);

    bone_number = first_mesh->mNumBones;
    // 每个骨骼影响了哪些顶点，权重是多少
    for (int i = 0; i < first_mesh->mNumBones; i++)
    {
        auto bone = first_mesh->mBones[i];

        float bone_mat[16] = {
                bone->mOffsetMatrix.a1, bone->mOffsetMatrix.a2, bone->mOffsetMatrix.a3, bone->mOffsetMatrix.a4,
                bone->mOffsetMatrix.b1, bone->mOffsetMatrix.b2, bone->mOffsetMatrix.b3, bone->mOffsetMatrix.b4,
                bone->mOffsetMatrix.c1, bone->mOffsetMatrix.c2, bone->mOffsetMatrix.c3, bone->mOffsetMatrix.c4,
                bone->mOffsetMatrix.d1, bone->mOffsetMatrix.d2, bone->mOffsetMatrix.d3, bone->mOffsetMatrix.d4
        };

        LOGI("SkinnedMesh bone %d, name = %s, mNumWeights = %d", i, bone->mName.C_Str(), bone->mNumWeights);

        bone_name_2_index[bone->mName.C_Str()] = i;
        bone_matrix.emplace_back(glm::make_mat4x4(bone_mat));

#ifdef SKINMESHED_DEBUG
        LOGI("SkinnedMesh %f, %f, %f, %f", bone->mOffsetMatrix.a1, bone->mOffsetMatrix.a2, bone->mOffsetMatrix.a3, bone->mOffsetMatrix.a4);
        LOGI("SkinnedMesh %f, %f, %f, %f", bone->mOffsetMatrix.b1, bone->mOffsetMatrix.b2, bone->mOffsetMatrix.b3, bone->mOffsetMatrix.b4);
        LOGI("SkinnedMesh %f, %f, %f, %f", bone->mOffsetMatrix.c1, bone->mOffsetMatrix.c2, bone->mOffsetMatrix.c3, bone->mOffsetMatrix.c4);
        LOGI("SkinnedMesh %f, %f, %f, %f", bone->mOffsetMatrix.d1, bone->mOffsetMatrix.d2, bone->mOffsetMatrix.d3, bone->mOffsetMatrix.d4);
#endif

        for (int j = 0; j < bone->mNumWeights; j++)
        {
            auto weight = bone->mWeights[j];
            vertex_weight[weight.mVertexId].emplace_back(i, weight.mWeight);
        }

    }

    for (int i = 0 ; i < vertexNumber; i++)
    {
        std::sort(vertex_weight[i].begin(), vertex_weight[i].end(), &BoneVertexData_sorter);
#ifdef SKINMESHED_DEBUG
        std::stringstream ss;
        for (int j = 0; j < vertex_weight[i].size(); j++)
        {
            ss << "(" << vertex_weight[i][j].bone_index << ", " << vertex_weight[i][j].bone_weight << "), ";
        }
        LOGI("SkinnedMesh %d %s", i, ss.str().c_str());
#endif

        bone_index_data[i * 4 + 0] = 0;
        bone_index_data[i * 4 + 1] = 0;
        bone_index_data[i * 4 + 2] = 0;
        bone_index_data[i * 4 + 3] = 0;

        bone_weight_data[i * 4 + 0] = 0.0f;
        bone_weight_data[i * 4 + 1] = 0.0f;
        bone_weight_data[i * 4 + 2] = 0.0f;
        bone_weight_data[i * 4 + 3] = 0.0f;

        for (int j = 0; j < 4 && j < vertex_weight[i].size(); j ++)
        {
            bone_index_data[i * 4 + j] = vertex_weight[i][j].bone_index;
            bone_weight_data[i * 4 + j] = vertex_weight[i][j].bone_weight;
        }
    }

    bone_matrix_list.resize(16 * bone_number);
    TravelBoneHierarchy(ai_scene->mRootNode, 0, glm::identity<glm::mat4x4>());
}

void SkinnedMeshModel::Update(double elapsed_time, double tot_time)
{
    TravelBoneHierarchy(ai_scene->mRootNode, tot_time, glm::identity<glm::mat4x4>());
}

float * SkinnedMeshModel::GetBoneTransformMatrixList()
{
    return bone_matrix_list.data();
}

int SkinnedMeshModel::GetBoneNumber()
{
    return bone_number;
}

void SkinnedMeshModel::TravelBoneHierarchy(aiNode* parent, double time, const glm::mat4x4 & parent_transform)
{
    std::string bone_name = parent->mName.C_Str();
    glm::mat4x4 bone_transform = glm::identity<glm::mat4x4>();
    int index = -1;

    double tot_time_in_ms = time * 1000.0f;
    auto anim = ai_scene->mAnimations[0];

    while (tot_time_in_ms > anim->mDuration)
    {
        tot_time_in_ms -= anim->mDuration;
    }

    if (bone_name_2_index.find(bone_name) != bone_name_2_index.end())
    {
        index = bone_name_2_index[bone_name];
        bone_transform = bone_matrix[index];

        /*
        auto node = anim->mChannels[index];

        aiVector3D pos;
        aiQuaternion rotation;
        aiVector3D scale;

        for (int i = 0; i < node->mNumPositionKeys - 1; i ++)
        {
            if (node->mPositionKeys[i].mTime < tot_time_in_ms && tot_time_in_ms < node->mPositionKeys[i + 1].mTime)
            {
                double t = (tot_time_in_ms - node->mPositionKeys[i].mTime) / (node->mPositionKeys[i + 1].mTime - node->mPositionKeys[i].mTime);
                pos = node->mPositionKeys[i].mValue * (float) t + node->mPositionKeys[i + 1].mValue * (float) (1 - t);
            }
        }

        for (int i = 0; i < node->mNumRotationKeys - 1; i ++)
        {
            if (node->mRotationKeys[i].mTime < tot_time_in_ms && tot_time_in_ms < node->mRotationKeys[i + 1].mTime)
            {
                double t = (tot_time_in_ms - node->mRotationKeys[i].mTime) / (node->mRotationKeys[i + 1].mTime - node->mRotationKeys[i].mTime);
                aiQuaternion::Interpolate(rotation, node->mRotationKeys[i].mValue, node->mRotationKeys[i + 1].mValue, t);
            }
        }

        for (int i = 0; i < node->mNumScalingKeys - 1; i ++)
        {
            if (node->mScalingKeys[i].mTime < tot_time_in_ms && tot_time_in_ms < node->mScalingKeys[i + 1].mTime)
            {
                double t = (tot_time_in_ms - node->mScalingKeys[i].mTime) / (node->mScalingKeys[i + 1].mTime - node->mScalingKeys[i].mTime);
                scale = node->mScalingKeys[i].mValue * (float) t + node->mScalingKeys[i + 1].mValue * (float) (1 - t);
            }
        }

        glm::quat bone_rotation;
        glm::mat4 translation_mat = glm::mat4(1.0f);
        translation_mat[3][0] = pos.x;
        translation_mat[3][1] = pos.y;
        translation_mat[3][2] = pos.z;

        glm::mat4 scale_mat = glm::mat4(1.0f);
        scale_mat[0][0] = scale.x;
        scale_mat[1][1] = scale.y;
        scale_mat[2][2] = scale.z;

        bone_transform = translation_mat * glm::mat4_cast(bone_rotation);

        LOGI("bone_name = %s, pos = (%f, %f, %f), rotation = (%f, %f, %f, %f), scale = (%f, %f, %f)",
                bone_name.c_str(), pos.x, pos.y, pos.z, rotation.x, rotation.y, rotation.z, rotation.w, scale.x, scale.y, scale.z);

        glm::vec3 glm_scale;
        glm::quat glm_rotation;
        glm::vec3 glm_translation;
        glm::vec3 glm_skew;
        glm::vec4 glm_perspective;
        glm::decompose(bone_transform, glm_scale, glm_rotation, glm_translation, glm_skew, glm_perspective);
        LOGI("bone_transform pos = (%f, %f, %f), rotation = (%f, %f, %f, %f), scale = (%f, %f, %f)",
                glm_translation.x, glm_translation.y, glm_translation.z,
                glm_rotation.x, glm_rotation.y, glm_rotation.z, glm_rotation.w,
                glm_scale.x, glm_scale.y, glm_scale.z);

        */
    }

    glm::mat4x4 real_transform = bone_transform * parent_transform;

    if (index >= 0) {
        float * mat_value = glm::value_ptr(real_transform);
        for (int i = 0; i < 16; i++) {
            bone_matrix_list[i + 16 * index] = mat_value[i];
        }
    }

    for (int j = 0; j < parent->mNumChildren; j ++)
    {
        TravelBoneHierarchy(parent->mChildren[j], time, real_transform);
    }
}